import torch.nn as nn
import torch.optim as optim
import torch.utils.data.distributed
import torchvision.models

from sort.data.DataHandler import DataHandler, Trainer

if __name__ == '__main__':

    # 训练目录和批次
    data_dir = 'I:\\Workspace\\Dataset\\data_inception\\train'
    batch_size = 16
    # device = "cpu"
    # 判断环境是CPU运行还是GPU
    device = (
        "cuda"
        if torch.cuda.is_available()
        else "mps"
        if torch.backends.mps.is_available()
        else "cpu"
    )
    print(device)

    # 加载数据集
    data_handler = DataHandler(data_dir, batch_size)
    train_dataset, val_dataset = data_handler.preprocess_data()
    train_loader, val_loader = data_handler.preprocess_data_loader()

    # 根据数据集确定类 和 类数量
    classes = train_dataset.dataset.classes
    num_classes = len(classes)

    # 定义ResNet-50模型
    # 初始化 model 进行网络下载
    # weights = torchvision.models.ResNet50_Weights.DEFAULT
    # model = torchvision.models.resnet50(weights=weights)

    # 本地加载
    model = torchvision.models.resnet50(weights=None).to(device)
    model.load_state_dict(torch.load('../../resource/resnet/resnet50-0676ba61.pth'))
    model = model.to(device)
    # 查看模型
    model.eval()

    num_features = model.fc.in_features
    model.fc = nn.Linear(num_features, num_classes)

    # 设置损失函数和优化器
    loss_function = nn.CrossEntropyLoss()
    # 创建优化器
    optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)

    # 训练结果查看 & 反向传播和优化
    trainer = Trainer(model, device, loss_function, optimizer)

    # 训练批次
    epochs = 5
    for t in range(epochs):
        print(f"Epoch {t + 1}\n-------------------------------")

        trainer.model = trainer.model.to(device)

        trainer.train(train_loader)
        trainer.test(val_loader)

    model_path = '../../model/ResNet50-ImageSort.pth'

    # 保存模型结构和权重
    # torch.save(model, model_path)

    # 或者只保存模型权重
    torch.save(model.state_dict(), model_path)
    print("Finished Training!")
